package middleware

import (
	"context"
	"fmt"
	"github.com/gofiber/fiber/v2"
	"github.com/google/uuid"
	"github.com/rs/zerolog/log"
	"shyxy-model-agent/core/config"
	"shyxy-model-agent/core/schema"
	"shyxy-model-agent/pkg/model"
)

type correlationIDKeyType string

const CorrelationIDKey correlationIDKeyType = "correlationID"

type RequestExtractor struct {
	applicationConfig *config.ApplicationConfig
	modelLoader       *model.ModelLoader
	dbLoader          *config.DBLoader
}

func NewRequestExtractor(dbLoader *config.DBLoader, modelLoader *model.ModelLoader, applicationConfig *config.ApplicationConfig) *RequestExtractor {
	return &RequestExtractor{
		applicationConfig: applicationConfig,
		modelLoader:       modelLoader,
		dbLoader:          dbLoader,
	}
}

const CONTEXT_LOCALS_KEY_MODEL_ID = "MODEL_ID"
const CONTEXT_LOCALS_KEY_MODEL_NAME = "MODEL_NAME"
const CONTEXT_LOCALS_KEY_DEF_MODEL_REQUEST = "DEF_MODEL_REQUEST"

func (re *RequestExtractor) BuildConstantDefaultModelIDMiddleware(defaultModelID string) fiber.Handler {
	return func(ctx *fiber.Ctx) error {
		re.setModelIDFromRequest(ctx)
		localModelID, ok := ctx.Locals(CONTEXT_LOCALS_KEY_MODEL_ID).(string)
		if !ok || localModelID == "" {
			ctx.Locals(CONTEXT_LOCALS_KEY_MODEL_ID, defaultModelID)
			log.Debug().Str("defaultModelID", defaultModelID).Msg("context local model name not found, setting to default")
		}
		return ctx.Next()
	}
}

func (re *RequestExtractor) setModelIDFromRequest(ctx *fiber.Ctx) {
	// local model id
	modelID, ok := ctx.Locals(CONTEXT_LOCALS_KEY_MODEL_ID).(string)
	if ok && modelID != "" {
		return
	}
	// http model id
	// GET /v1/models/:modelID
	modelID = ctx.Params("model")
	// GET /v1/models?modelID=
	if (modelID == "") && ctx.Query("model") != "" {
		modelID = ctx.Query("model")
	}

	log.Debug().Str("model", modelID).Msg("context local model name found")
	ctx.Locals(CONTEXT_LOCALS_KEY_MODEL_ID, modelID)
}

func (re *RequestExtractor) SetModelAndConfig(initializer func() schema.DefModelRequest) fiber.Handler {
	return func(ctx *fiber.Ctx) error {
		input := initializer()
		if input == nil {
			return fmt.Errorf("unable to initialize body")
		}
		if err := ctx.BodyParser(input); err != nil {
			return fmt.Errorf("failed parsing request body: %w", err)
		}

		// If this request doesn't have an associated model name, fetch it from earlier in the middleware chain
		if input.ModelName(nil) == "" {
			localModelName, ok := ctx.Locals(CONTEXT_LOCALS_KEY_MODEL_NAME).(string)
			if ok && localModelName != "" {
				log.Debug().Str("context localModelName", localModelName).Msg("overriding empty model name in request body with value found earlier in middleware chain")
				input.ModelName(&localModelName)
			}
		}

		log.Debug().Msgf("input: %v", input)
		ctx.Locals(CONTEXT_LOCALS_KEY_DEF_MODEL_REQUEST, input)
		return ctx.Next()
	}
}

func (re *RequestExtractor) SetOpenAIRequest(ctx *fiber.Ctx) error {
	input, ok := ctx.Locals(CONTEXT_LOCALS_KEY_DEF_MODEL_REQUEST).(*schema.OpenAIRequest)
	if !ok || input.Model == "" {
		return fiber.ErrBadRequest
	}

	// Extract or generate the correlation ID
	correlationID := ctx.Get("X-Correlation-ID", uuid.New().String())
	ctx.Set("X-Correlation-ID", correlationID)
	c1, cancel := context.WithCancel(context.Background())
	// Add the correlation ID to the new context
	ctxWithCorrelationID := context.WithValue(c1, CorrelationIDKey, correlationID)

	input.Context = ctxWithCorrelationID
	input.Cancel = cancel

	ctx.Locals(CONTEXT_LOCALS_KEY_DEF_MODEL_REQUEST, input)

	return ctx.Next()
}
